import os
import m3u8  
import requests  
import concurrent.futures 
# import transcode  # ffmpeg转码
import shutil  
from urllib.parse import urlparse 
from datetime import datetime  
import hashlib

# ------- 工具 ------- #
def md5_hash(input_string):  
    md5 = hashlib.md5()  
    md5.update(input_string.encode('utf-8'))  
    return md5.hexdigest()

def printTime(str) :
    # 获取当前时间  
    current_time = datetime.now()  
    # 格式化时间  
    formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")  

    # 输出带有时间的信息  
    print(str + " " + formatted_time)

def getUrlLastPath(url):
    result = urlparse(url)  
    return result.path.split('/')[-1]


class M3u8Downloader: 
    def __init__(self, callback, maxRepeatTimes=3, downloadPath='./downloads'):
        self.downloadedCount = 0      # 已下载的数量
        self.totalCount = 0       # 总下载的文件数量：ts和key的总和
        self.failedCount = 0          # 下载失败的数量
        self.repeatTimesDic = {}      # 失败的文件，已重试的次数
        self.maxRepeatTimes = maxRepeatTimes # 最大的重试次数
        self.downloadPath = downloadPath     # mp4视频文件的最终存储目录
        self.defaultDownloadPath = downloadPath
        self.fileName = ''
        self.callback = callback


    # 解析得到m3u8实例
    def parseM3u8(self, m3u8Url):
        headers = {"userAgent":"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Safari/605.1.15",
        "Referer":"http://456dui.com/"}
        playlist = m3u8.load(m3u8Url, headers=headers)
        if playlist.is_variant:
            playlists = playlist.playlists
            if len(playlists) > 0 :
                subPlaylist = playlists[0]   # 默认取第一个，有多个的情况，其实可以优化（应该是选择分辨率高的）
                #subM3u8Url = subPlaylist.base_uri + subPlaylist.uri
                subM3u8Url = subPlaylist.absolute_uri
                return self.parseM3u8(subM3u8Url)
            else :
                return None
        else :
            return playlist


    # 参数校验
    def validateParameter(self, p, isUrl) :
        # 判断是否是字符串类型
        isStr = isinstance(p, str)
        if isStr :
            if isUrl :
                hasHttp = p.startswith('http://') or p.startswith('https://') 
                if len(p) > 15 and hasHttp :
                    return True 
            else :
                if len(p) > 3 :
                    return True
        return False


    # 从m3u8中，得到ts和key的链接，以及文件的存储全路径
    #def getTsAndKeyUrls(self, playlist): 
    def getFileUrlsAndPaths(self, playlist): 
        allUrls = []
        allFilePaths = []
        tempPath = self.tempPath(self.downloadPath, self.fileName)
        for seg in playlist.segments:
            if seg == None:
                break
            allUrls.append(seg.absolute_uri)
            fn = md5_hash(getUrlLastPath(seg.absolute_uri))
            na = tempPath + '/' + fn#getUrlLastPath(seg.absolute_uri)
            allFilePaths.append(na)

        for key in playlist.keys:
            if key == None:
                break
            allUrls.append(key.absolute_uri)
            # 将url中获取的文件名md5一下
            fn = md5_hash(getUrlLastPath(key.absolute_uri))
            na = tempPath + '/' + fn #getUrlLastPath(key.absolute_uri)
            allFilePaths.append(na)

        return allUrls, allFilePaths
        #return tsUrls, tsNames, keyUrls, keyNames

    # 创建本地m3u8文件，用于ffmpeg的ts合并以及转码mp4
    def createNativeM3u8File(self, playlist, m3u8Path): 

        for seg in playlist.segments:
            if seg == None:
                break
            seg.uri = md5_hash(getUrlLastPath(seg.absolute_uri))

        for key in playlist.keys:
            if key == None:
                break
            key.uri = md5_hash(getUrlLastPath(key.absolute_uri))

        playlist.dump(m3u8Path)

    # 重试下载
    def retryDownload_file(self, url, filePath):
        # 若下载失败，则开始重试
        allKeys = self.repeatTimesDic.keys()
        retryTimes = 0
        if url in allKeys :
            retryTimes = self.repeatTimesDic[url]

        print(f"无法下载文件 {fn}: {response.status_code}")  
        if retryTimes >= self.maxRepeatTimes : 
            print(f'文件 {fn} 已达到最大下载次数 {self.maxRepeatTimes} 次')
            print(f'文件url:{url}')
            self.failedCount += 1

        else :
            retryTimes = retryTimes + 1
            self.repeatTimesDic[url] = retryTimes
            print(f'开始第{retryTimes}次重试')
            self.download_file(url, filePath)
  
    # 下载单个文件：url文件的地址，filepath文件的存储全路径
    def download_file(self, url, filePath):  

        fn = getUrlLastPath(filePath)
        
        if os.path.exists(filePath):

            self.downloadedCount += 1  
            cnt = f"({self.downloadedCount}/{self.totalCount})"
            print(f"文件 {fn} 已存在{cnt}")

            return

        try:
            headers = {"userAgent":"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Safari/605.1.15",
        "Referer":"http://456dui.com/"}
            response = requests.get(url, headers=headers, verify=False)  
            if response.status_code == 200:  
                with open(filePath, 'wb') as file:  
                    file.write(response.content)

                self.downloadedCount += 1  
                cnt = f"({self.downloadedCount}/{self.totalCount})"
                print(f"文件 {fn} 下载完成{cnt}") 
            else:  
                self.retryDownload_file(url, filePath)
        except requests.exceptions.RequestException as e: 
            print(f"An error occurred异常: {e}")  
            self.retryDownload_file(url, filePath)
                

    # 多线程下载
    def download_files_concurrently(self, url_list, filename_list):  
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:  
            # executor.map(lambda x: self.download_file(*x), zip(url_list, filename_list))   
            # 将 url_list 和 filename_list 中的对应元素作为参数传递给 download_file 方法  
            tasks = [executor.submit(self.download_file, url, filename) for url, filename in zip(url_list, filename_list)]  
              
            # 等待所有任务完成  
            for future in concurrent.futures.as_completed(tasks):  
                # Future.result() 会阻塞并等待任务完成，并返回任务结果（这里是 None）  
                result = future.result()  
                if result is not None:  # 如果结果不为 None，则说明任务完成了  
                    print("Task completed:", result)  
              
            # 所有任务都已完成  
            printTime("All tasks are done.")
            return 1
            # self.downloadComplete()

    # 临时目录，存放m3u8\ts\key等文件
    def tempPath(self, downloadPath, fileName) :
        tempPath = downloadPath + '/' + fileName

        # 若不存在，创建临时目录
        if os.path.isdir(tempPath) == False :
            os.makedirs(tempPath, exist_ok=True)

        return tempPath

    # m3u8文件存储的路径
    def m3u8Path(self, downloadPath, fileName) : 
        return self.tempPath(downloadPath, fileName) + '/index.m3u8'

    # mp4的全路径
    def mp4Path(self, downloadPath, fileName) :
        # 带有拓展名的文件名
        fileAllName = fileName.replace(".mp4", "")+".mp4"

        # 文件输出路径
        mp4Path = downloadPath + '/' +fileAllName
        return mp4Path

    def downloadM3u8(self, m3u8Url, fileName, downloadPath=''):

        printTime("开始下载m3u8")
        # return 0

        self.downloadPath = downloadPath
        if len(self.downloadPath) < 3 :
            self.downloadPath = self.defaultDownloadPath

        if self.validateParameter(m3u8Url, True) == False :
            print('m3u8的url不正确，请重新输入')
            return 
        if self.validateParameter(fileName, False) == False :
            print('请输入正确的文件名')
            return 

        self.fileName = fileName
        
        # 文件输出路径
        mp4Path = self.mp4Path(self.downloadPath, self.fileName)

        # 检查mp4文件是否已存在
        if os.path.exists(mp4Path) and os.path.isdir(mp4Path) == False:  
            print(f'[{fileName}] 文件已存在！！！')
            return

        # 开始解析m3u8，并获取ts和key的url等
        printTime('开始获取m3u8的内容')
        print(f'm3u8URL: {m3u8Url}')
        playlist = self.parseM3u8(m3u8Url)

        # 得到ts和key的urls 以及 存储路径
        tsUrls, tsNames = self.getFileUrlsAndPaths(playlist)
        
        printTime('解析m3u8完成，得到所有的文件url')

        # 生成本地m3u8文件，用于ffmpeg合并ts并转码为mp4
        m3u8Path = self.m3u8Path(self.downloadPath, self.fileName)
        self.createNativeM3u8File(playlist, m3u8Path)


        self.downloadedCount = 0
        self.failedCount = 0
        self.totalCount = len(tsUrls)


        # 下载ts和key
        if len(tsUrls) > 0 and len(tsNames) > 0:

            printTime('开始下载所有文件(ts和key)')
            self.download_files_concurrently(tsUrls,tsNames)
            return 1
        else:
            printTime(f'文件 {fileName} 下载失败!')
            return 0

    # 文件全部下载完成，开始转码
    def downloadComplete(self) :
        if self.failedCount > 0 :
            printTime(f'下载失败！有{self.failedCount}个文件未能下载，请重试！！')
        else :
            printTime(f'所有文件已下载完成！！！下面开始转码！！！')

        # ts文件全部下载完成，开始转码 
        m3u8Path = self.m3u8Path(self.downloadPath, self.fileName)
        mp4Path = self.mp4Path(self.downloadPath, self.fileName)
        tempPath = self.tempPath(self.downloadPath, self.fileName)
        result = transcode.m3u8ToMp4(m3u8Path, mp4Path)
        if result :
            printTime(f'文件 [{self.fileName}] 转码完成！')
            self.callback(mp4Path)
            # 清空所有的临时文件
            if os.path.isdir(tempPath):  
                shutil.rmtree(tempPath) 
                printTime("已清空临时文件！！")  
        else :
            printTime(f'文件 [{self.fileName}] 转码失败！')
            self.callback(None)

    # m3u8转mp4的命令
    def getFFmpegCommand(self) :
        m3u8Path = self.m3u8Path(self.downloadPath, self.fileName)
        mp4Path = self.mp4Path(self.downloadPath, self.fileName)
        ffmpeg_command = f"ffmpeg -allowed_extensions ALL -protocol_whitelist \"file,http,crypto,tcp\" -i {m3u8Path} -c copy {mp4Path}"  
        

        ffmpeg_command = "ffmpeg -allowed_extensions ALL -protocol_whitelist 'file,http,crypto,tcp' -i ./m3u8dlDownloads/bsws/index.m3u8 -c copy ./m3u8dlDownloads/bsws.mp4"

        return ffmpeg_command

        # return (m3u8Path, mp4Path)

def m3u8dl(m3u8url, filename) :
    md = M3u8Downloader()
    ret = md.downloadM3u8(m3u8url, filename)
    return ret 

# # 示例用法
# md = M3u8Downloader()
# ul = 'https://hls.cntv.cdn20.com/asp/hls/1200/0303000a/3/default/1049a60b8b914d4ab7066c568ca616fd/1200.m3u8'
# na = '小孩子大梦想'
# md.downloadM3u8(ul, na)



